import torch
import torchvision.transforms as transforms
import numpy as np
import argparse
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import StepLR
from dataset import MyDataset
from model import ContactAngle
from config import *
from utils import setup_seed, Save_data, Loss_function

parser = argparse.ArgumentParser('Prediction Model for Contact Angle')
parser.add_argument('--it', default=50, type=int, help='训练次数')
parser.add_argument('--lr', default=1e-4, type=float, help='学习率')
parser.add_argument('--ps', default=224, type=int, help='图像大小')
parser.add_argument('--bs', default=32, type=int, help='迷你批次大小')
parser.add_argument('--ratio', default=0.8, type=float, help='训练集/(训练集+测试集)')
parser.add_argument('--resume', default=False, type=str, help='是否从检查点加载权重参数')
parser.add_argument('--ss', default=10, type=int, help='步长')
parser.add_argument('--gm', default=0.1, type=float, help='gamma值')
parser.add_argument('--wd', default=5e-4, type=float, help='权重衰减率')
parser.add_argument('--sd', default=100, type=int, help='随机种子参数')
parser.add_argument('--out', default=1020, type=int, help='图像提取网络输出向量维度')
parser.add_argument('--pandas', default=True, type=str, help='是否保存预测误差值')
parser.add_argument('--loss', default='L1', type=str, help='L1, L1SMOOTH, MSE')
parser.add_argument('--mn', default='Resnet34', type=str, help='Densnet121, Resnet18, Resnet34, Resnet50, Vgg16, '
                                                               'Resnet101,Vgg11, Alexnet, Vit_small_patch16_224, '
                                                               'Vit_base_patch16_224, Vit_base_patch32_224, '
                                                               'Vit_large_patch16_224, Vit_large_patch32_224, '
                                                               'Vit_base_resnet26d_224, Vit_base_resnet50d_224, '
                                                               'Vit_small_resnet26d_224')
args = parser.parse_args()

start_epoch = 0
best_loss = float('inf')
setup_seed(args.sd)

# 模型在CPU或者GPU上运行
if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

# 表面沟槽微结构图像预处理
transform = transforms.Compose([
    transforms.Resize((args.ps, args.ps)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

# 数据集加载
full_dataset = MyDataset(img_path=img_root, txt_path=txt_root, transform=transform)
train_size = int(args.ratio * len(full_dataset))
test_size = len(full_dataset) - train_size
train_data, test_data = torch.utils.data.random_split(full_dataset, [train_size, test_size])

train_data_loader = DataLoader(dataset=train_data, batch_size=args.bs, shuffle=True)
test_data_loader = DataLoader(dataset=test_data, batch_size=args.bs, shuffle=True)

# 接触角预测模型
model = ContactAngle(model_name=args.mn, out_puts=args.out)

# Adam自适应优化器
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
# 固定步长衰减策略
scheduler = StepLR(optimizer, step_size=args.ss, gamma=args.gm)

# 加载模型原始参数或者检查点参数
if args.resume:
    print('Resuming from checkpoint')
    checkpoint = torch.load(checkpoint_root + r'\checkpoint-%s.pth' % args.mn)
    model.load_state_dict(checkpoint['model'])
    start_epoch = checkpoint['epoch']
    best_loss = checkpoint['loss']
    optimizer.load_state_dict(checkpoint['optimizer'])
    scheduler.load_state_dict(checkpoint['scheduler'])
model = model.to(device)

# 模型损失函数
Loss = Loss_function[args.loss]

# 模型训练过程
def train(epoch):
    epoch_data = []
    model.train()
    print('\nTraining Epoch:%d' % (epoch + 1))
    train_loss = 0
    for step, (img, txt, label) in enumerate(train_data_loader):
        img = img.to(device)
        txt = txt.to(device)
        label = label.to(device)
        outputs = model(img, txt)
        optimizer.zero_grad()
        loss = Loss(outputs, label)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        if args.pandas:
            epoch_data.append(train_loss/(step + 1))
        print('Step:[%d/%d], Loss:%.4f, Train_loss:%.4f' % (
            step + 1, len(train_data_loader), loss.item(), train_loss / (step + 1)))
    scheduler.step()
    return epoch_data

# 模型测试过程
def test(epoch):
    model.eval()
    print('\nTesting Epoch:%d' % (epoch + 1))
    with torch.no_grad():
        test_loss = 0
        for step, (img, txt, label) in enumerate(test_data_loader):
            img = img.to(device)
            txt = txt.to(device)
            label = label.to(device)
            outputs = model(img, txt)
            loss = Loss(outputs, label)
            test_loss += loss.item()

        global best_loss
        test_loss /= len(test_data_loader)
        print('The Testing Loss is:%.4f' % test_loss)
        if test_loss < best_loss:
            print('Saving Model..........')
            state = {
                'epoch': epoch,
                'optimizer': optimizer.state_dict(),
                'scheduler': scheduler.state_dict(),
                'model': model.state_dict(),
                'loss': test_loss
            }
            if not os.path.isdir(checkpoint_root):
                os.makedirs(checkpoint_root)
            torch.save(state, checkpoint_root + r'\checkpoint-%s.pth' % args.mn)
            best_loss = test_loss
            if not os.path.isdir(weight_root):
                os.makedirs(weight_root)
            torch.save(model.state_dict(), weight_root + r'\model-%s.pth' % args.mn)


if __name__ == '__main__':
    save_data = []
    for i in range(start_epoch, start_epoch + args.it):
        step_data = train(i)
        save_data.append(step_data)
        test(i)

    # 保存预测损失值
    if args.pandas:
        cacul_data = np.array(save_data).reshape(-1)
        Save_data(data=cacul_data, root=csv_root + r'\%s.csv' % args.mn)
